{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 吉布斯态的制备"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概览\n",
    "\n",
    "在本案例中，我们将展示如何通过 Paddle Quantum 训练量子神经网络（quantum neural network, QNN）来制备量子吉布斯态。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 背景\n",
    "\n",
    "量子计算中的前沿方向包含量子机器学习和量子优化，在这两个方向中，特定量子态的制备是非常重要的问题。特别的，吉布斯态（Gibbs state）的制备是实现诸多量子算法所必须的步骤并且广泛应用于：\n",
    "\n",
    "- 量子机器学习中受限波尔兹曼机的学习 [1]\n",
    "- 解决凸优化和半正定规划等优化问题 [2]\n",
    "- 组合优化问题 [3]\n",
    "\n",
    "具体的吉布斯态定义如下：给定一个 $n$ 量子比特的哈密顿量 $H$（一般来说这是一个$2^n\\times2^n$的厄米矩阵），其在温度 $T$ 下的吉布斯态为 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\rho_G = \\frac{{{e^{ - \\beta H}}}}{{\\text{tr}({e^{ - \\beta H}})}},\n",
    "\\tag{1}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其中 ${e^{ - \\beta H}}$ 是矩阵 $ - \\beta H$ 的矩阵指数，$\\beta  = \\frac{1}{{kT}}$ 是系统的逆温度参数，$T$ 是温度参数，$k$ 是玻尔兹曼常数 (这篇教程中我们取 $k = 1$)。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paddle Quantum 实现"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先通过下面几行代码引入必要的 library 和 package。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:55:59.838299Z",
     "start_time": "2021-04-30T08:55:57.450922Z"
    }
   },
   "outputs": [],
   "source": [
    "import scipy\n",
    "import paddle\n",
    "from numpy import trace as np_trace\n",
    "import paddle_quantum as pq\n",
    "from paddle_quantum.ansatz import Circuit\n",
    "from paddle_quantum.state import zero_state\n",
    "from paddle_quantum.qinfo import state_fidelity, partial_trace, pauli_str_to_matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "作为一个上手的例子，这里我们考虑一个 3 量子比特的哈密顿量及其吉布斯态：\n",
    "\n",
    "$$\n",
    "H = -Z \\otimes Z \\otimes I - I \\otimes Z \\otimes Z - Z \\otimes I \\otimes Z, \\quad I=\\left [\n",
    "\\begin{matrix}\n",
    "1 & 0  \\\\\n",
    "0 & 1  \\\\\n",
    "\\end{matrix} \n",
    "\\right ], \\quad \n",
    "Z=\\left [\n",
    "\\begin{matrix}\n",
    "1 & 0  \\\\\n",
    "0 & -1  \\\\\n",
    "\\end{matrix} \n",
    "\\right ].\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "这个例子中，我们将逆温度参数设置为 $\\beta = 1.5$。此外，为了方便测试结果，我们按照定义提前生成好了理想情况的吉布斯态 $\\rho_G$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:55:59.855633Z",
     "start_time": "2021-04-30T08:55:59.841241Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 4                               # 量子神经网络的宽度\n",
    "N_SYS_B = 3                         # 用于生成吉布斯态的子系统B的量子比特数   \n",
    "SEED = 16                           # 固定随机种子\n",
    "beta = 1.5                          # 设置逆温度参数 beta\n",
    "pq.set_backend('density_matrix')    # 设置密度矩阵后端\n",
    "pq.set_dtype('complex128')          # 设置计算精度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:55:59.935612Z",
     "start_time": "2021-04-30T08:55:59.866942Z"
    }
   },
   "outputs": [],
   "source": [
    "# 生成用泡利字符串表示的特定的哈密顿量\n",
    "H = [[-1.0, 'z0,z1'], [-1.0, 'z1,z2'], [-1.0, 'z0,z2']]\n",
    "\n",
    "# 生成哈密顿量的矩阵信息\n",
    "hamiltonian = pauli_str_to_matrix(H, N_SYS_B).numpy()\n",
    "\n",
    "# 生成理想情况下的目标吉布斯态 rho\n",
    "rho_G = scipy.linalg.expm(-1 * beta * hamiltonian) / np_trace(scipy.linalg.expm(-1 * beta * hamiltonian))\n",
    "\n",
    "# 设置成 Paddle quantum 所支持的数据类型\n",
    "hamiltonian = hamiltonian.astype(\"complex128\")\n",
    "rho_G = paddle.to_tensor(rho_G, dtype=\"complex128\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 搭建量子神经网络"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 在这个案例中，我们将通过训练量子神经网络（也可以理解为参数化量子电路）来制备吉布斯态。这里，我们提供一个简单的 4 量子比特的量子电路如下：\n",
    "\n",
    "  ![Ugibbs.jpg](https://release-data.cdn.bcebos.com/PIC%2FUgibbs.jpg)\n",
    "\n",
    "- 我们需要预设一些电路的参数，比如电路有 4 个量子比特，其中第 1 个量子比特是辅助系统，第 2-4 个量子比特是用以产生吉布斯态的子系统。\n",
    "\n",
    "- 初始化其中的变量参数，$\\theta$ 代表我们量子神经网络中的参数组成的向量。\n",
    "         \n",
    "\n",
    "接下来我们根据上图中的电路设计，通过 Paddle Quantum 的 `Circuit` 类和内置的 `real_entangled_layer` 电路模板来高效搭建量子神经网络。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:01.029894Z",
     "start_time": "2021-04-30T08:56:01.022404Z"
    }
   },
   "outputs": [],
   "source": [
    "def U_theta(num_qubits: int, depth: int) -> Circuit:\n",
    "    \"\"\"\n",
    "    量子神经网络\n",
    "    \"\"\"\n",
    "    cir = Circuit(num_qubits)\n",
    "    \n",
    "    # 内置的 {R_y + CNOT} 电路模板\n",
    "    cir.real_entangled_layer(depth=depth)\n",
    "    \n",
    "    # 铺上最后一层 R_y 旋转门\n",
    "    cir.ry()\n",
    "    \n",
    "    return cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 配置训练模型——损失函数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 现在我们已经有了数据和量子神经网络的架构，我们将进一步定义合适的训练参数、模型和损失函数来达到我们的目标。\n",
    "\n",
    "- 具体的我们参考的是论文 [4] 中的方法，核心思想是**利用吉布斯态达到了最小自由能**的性质。\n",
    "\n",
    "- 通过作用量子神经网络 $U(\\theta)$ 在初始态上，我们可以得到输出态 $\\left| {\\psi \\left( {\\bf{\\theta }} \\right)} \\right\\rangle $，其在第 2-4 个量子比特的态记为 $\\rho_B(\\theta)$。\n",
    "\n",
    "- 设置训练模型中的的损失函数。在吉布斯态学习中，我们利用冯诺依曼熵函数的截断来进行自由能的估计，相应的损失函数参考 [4] 可以设为 $loss= {L_1} + {L_2} + {L_3}$，其中 \n",
    "\n",
    "$$\n",
    "{L_1}= \\text{tr}(H\\rho_B), \\quad {L_2} = 2{\\beta^{-1}}{\\text{tr}}(\\rho_B^2), \\quad L_3 = - {\\beta ^{ - 1}}\\big(\\text{tr}(\\rho_B^3) + 3\\big)/2.\n",
    "\\tag{3}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:03.701245Z",
     "start_time": "2021-04-30T08:56:03.691389Z"
    }
   },
   "outputs": [],
   "source": [
    " # 定义损失函数\n",
    "def loss_func(cir: Circuit, Hamiltonian: paddle.Tensor, N_SYS_B: int) -> paddle.Tensor:\n",
    "    # 施加量子神经网络\n",
    "    rho_AB = cir(zero_state(N))\n",
    "    \n",
    "    # 计算偏迹 partial trace 来获得子系统B所处的量子态 rho_B\n",
    "    rho_B = partial_trace(rho_AB.data, 2 ** (N - N_SYS_B), 2 ** (N_SYS_B), 1)\n",
    "    \n",
    "    # 计算三个子损失函数\n",
    "    rho_B_squre = rho_B @ rho_B\n",
    "    loss1 = paddle.real(paddle.trace(rho_B @ Hamiltonian))\n",
    "    loss2 = paddle.real(paddle.trace(rho_B_squre)) * 2 / beta\n",
    "    loss3 = -(paddle.real(paddle.trace(rho_B_squre @ rho_B)) + 3) / (2 * beta)\n",
    "    \n",
    "    # 最终的损失函数\n",
    "    loss = loss1 + loss2 + loss3  \n",
    "    \n",
    "    return loss, rho_B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 配置训练模型——模型参数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在进行量子神经网络的训练之前，我们还需要进行一些训练的超参数设置，主要是学习速率（learning rate, LR）、迭代次数（iteration, ITR）和量子神经网络计算模块的深度（depth, D）。这里我们设定学习速率为 0.5，迭代次数为 50 次。读者不妨自行调整来直观感受下超参数调整对训练效果的影响。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:05.260360Z",
     "start_time": "2021-04-30T08:56:05.251628Z"
    }
   },
   "outputs": [],
   "source": [
    "ITR = 50 # 设置训练的总迭代次数\n",
    "LR = 0.5 # 设置学习速率\n",
    "D = 1    # 设置量子神经网络中重复计算模块的深度 Depth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 进行训练"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 当训练模型的各项参数都设置完成后，我们将数据转化为 Paddle 中的张量，进而进行量子神经网络的训练。\n",
    "- 训练过程中我们用的是 [Adam Optimizer](https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/optimizer/adam/Adam_cn.html)，也可以调用 Paddle 中提供的其他优化器。\n",
    "- 我们将训练过程中的结果依次输出。\n",
    "- 特别地，我们依次输出了我们学习到的量子态 $\\rho_B(\\theta)$ 与吉布斯态 $\\rho_G$ 的保真度，保真度越高说明QNN输出的态越接近于吉布斯态。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:56:19.753228Z",
     "start_time": "2021-04-30T08:56:15.379159Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 10 loss: -3.1085 fid: 0.9241\n",
      "iter: 20 loss: -3.3375 fid: 0.9799\n",
      "iter: 30 loss: -3.3692 fid: 0.9897\n",
      "iter: 40 loss: -3.3990 fid: 0.9929\n",
      "iter: 50 loss: -3.4133 fid: 0.9959\n",
      "\n",
      "训练后的电路:  \n",
      "--Ry(6.290)----*--------------x----Ry(0.747)--\n",
      "               |              |               \n",
      "--Ry(4.745)----x----*---------|----Ry(6.249)--\n",
      "                    |         |               \n",
      "--Ry(-0.01)---------x----*----|----Ry(-0.05)--\n",
      "                         |    |               \n",
      "--Ry(0.017)--------------x----*----Ry(6.310)--\n",
      "                                              \n"
     ]
    }
   ],
   "source": [
    "paddle.seed(SEED)\n",
    "    \n",
    "# 我们需要将 Numpy array 转换成 Paddle 中支持的 Tensor\n",
    "H = paddle.to_tensor(hamiltonian)\n",
    "\n",
    "# 确定网络的参数维度\n",
    "circuit = U_theta(N, D)\n",
    "\n",
    "# 一般来说，我们利用 Adam 优化器来获得相对好的收敛，\n",
    "# 当然你可以改成 SGD 或者是 RMS prop.\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=circuit.parameters())\n",
    "\n",
    "# 优化循环\n",
    "for itr in range(1, ITR + 1):\n",
    "        \n",
    "    # 前向传播计算损失函数并返回生成的量子态 rho_B\n",
    "    loss, rho_B = loss_func(circuit, H, N_SYS_B)\n",
    "        \n",
    "    # 反向传播极小化损失函数\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    opt.clear_grad()\n",
    "\n",
    "    # 转换成 Numpy array 用以计算量子态的保真度 F(rho_B, rho_G)\n",
    "    fid = state_fidelity(rho_B, rho_G)\n",
    "\n",
    "    # 打印训练结果\n",
    "    if itr % 10 == 0:\n",
    "        print('iter:', itr, 'loss:', '%.4f' % loss.numpy(), 'fid:', '%.4f' % fid.numpy())\n",
    "    if itr == ITR:\n",
    "        print(\"\\n训练后的电路:  \")\n",
    "        print(circuit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 总结"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "根据上面训练得到的结果，通过大概 50 次迭代，我们就能达到高于 99.5% 保真度的高精度吉布斯态，高效并精确地完成了吉布斯态的制备。我们可以通过 print 函数来输出学习到的量子神经网络的参数和它的输出态。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_______\n",
    "\n",
    "## 参考文献"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[1] Kieferová, M. & Wiebe, N. Tomography and generative training with quantum Boltzmann machines. [Phys. Rev. A 96, 062327 (2017).](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.96.062327)\n",
    "\n",
    "[2] Brandao, F. G. S. L. & Svore, K. M. Quantum Speed-Ups for Solving Semidefinite Programs. [in 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) 415–426 (IEEE, 2017). ](https://ieeexplore.ieee.org/abstract/document/8104077)\n",
    "\n",
    "[3] Somma, R. D., Boixo, S., Barnum, H. & Knill, E. Quantum Simulations of Classical Annealing Processes. [Phys. Rev. Lett. 101, 130504 (2008).](https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.101.130504)\n",
    "\n",
    "[4] Wang, Y., Li, G. & Wang, X. Variational quantum Gibbs state preparation with a truncated Taylor series. [Phys. Rev. A 16, 054035 (2021).](https://journals.aps.org/prapplied/abstract/10.1103/PhysRevApplied.16.054035)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('py3.7_pq2.2.1')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "4e4e2eb86ad73936e915e7c7629a18a8ca06348106cf3e66676b9578cb1a47dd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
